#!/usr/bin/python -S
"""
streams.py

Various streams used in xmap.
"""

__author__ = 'Andy Chu'


import os
import sys

from torn import base

import tnet

from util import log


# States for netstring parser
_NEED_LENGTH = 0   # initial state, we don't know how many bytes we need
_HAVE_LENGTH = 1   # we know how many bytes we need, but haven't gotten them yet


class TnetValueStream(base.EventEmitter):
  """Wraps a readable file-like emitter.

  It registers itself as a listener for 'data', and then emits complete 'value'
  events.  The 'close' event is echoed.
  """
  EMITS = [
      'chunk',  # a "framed" tnet value.  This is a string which is NOT decoded.
      'close',  # no more 'value' events are coming
      ]

  def __init__(self, in_stream, event_loop=None):
    """
    Args:
      in_stream: An event emitter that emits 'data' and 'close' events.
      'data' events
    """
    base.EventEmitter.__init__(self, event_loop=event_loop)

    self.buf = ''   # incomplete tnet string
    self.state = _NEED_LENGTH
    self.bytes_needed = -1  # number of bytes needed to emit

    in_stream.on('data', self.on_read)
    in_stream.on('close', self.on_close)
    self.in_stream = in_stream

  def pause(self):
    """Stop emitting event.
    
    Implemented by pausing the underlying stream.
    """
    log('TnetValueStream PAUSE')
    self.in_stream.pause()

  def resume(self):
    self.in_stream.resume()

  def _emit_and_reset(self):
    """Emit a value, and recalculate internal state."""
    #print 'EMITTING', repr(data[:bytes_needed])
    log('TnetValueStream emitting %d bytes', self.bytes_needed)
    self.emit('chunk', self.buf[:self.bytes_needed])
    self.buf = self.buf[self.bytes_needed:]

    # See if we have another value
    # TODO: really we should test for digits first, here and below
    i = self.buf.find(':', 0, 10)
    if i != -1:
      length_str = self.buf[:i]
      rest = self.buf[i+1:]
      length = int(length_str)
      self.bytes_needed = len(length_str) + 1 + length + 1
      self.state = _HAVE_LENGTH
    else:
      self.state = _NEED_LENGTH
      # TODO: This causes a bug!  I think there may be a bug here.
      #self.bytes_needed = -1  # not technically needed, but clearer

  def on_read(self, data):
    log('TnetValueStream on read %r  from state: %s', data, self.state)

    # append data to existing buffer
    self.buf += data

    if self.state == _NEED_LENGTH:
      # only search 10 chars, since it can't be big
      i = self.buf.find(':', 0, 10)
      if i == -1:
        return  # got more bytes, but still in the _NEED_LENGTH state

      # Process the length, and emit if we need to

      # self.buf is either EMPTY, or contains a PARTIAL length
      length_str = self.buf[:i]
      rest = self.buf[i+1:]
      #print repr(length_str), repr(rest)

      length = int(length_str)
      self.bytes_needed = len(length_str) + 1 + length + 1
      self.state = _HAVE_LENGTH

      # If we have enough bytes to emit, emit it, and reset the buffer in the
      # state.
      while len(self.buf) >= self.bytes_needed:
        self._emit_and_reset()

    elif self.state == _HAVE_LENGTH:
      # See if we can emit an event
      while len(self.buf) >= self.bytes_needed:
        self._emit_and_reset()

    else:
      raise AssertionError('Invalid state')

    log('TnetValueStream buf %r  to state: %s', self.buf, self.state)
    #print

  def on_close(self):
    log('****** TNET CLOSE')
    # Just propagate it outward
    self.emit('close')


class ValueStream(base.EventEmitter):
  EMITS = ['value', 'close']


class ChunkStream(base.EventEmitter):
  """
  This could be ByteStream, except the chunk boundaries matter.  Not valid to
  "reframe" chunks.
  """
  EMITS = ['chunk', 'close']


# How would a sentinel be represented on a inter-process thread?  This is just
# CLOSE.
SENTINEL = {}


class DiskThread(object):
  """Writes many logs to disk.

  event stream -> channel ->
  event stream -> channel -> multiplexer -> thread -> dmux -> disk
  """

  def __init__(self, in_channel, base_dir):
    """
    Args:
     dir: 
    """
    self.base_dir = base_dir
    self.in_channel = in_channel
    # name -> open file handle?
    self.handles = {}

  def _file(self, name):
    if name not in self.handles:
      path = os.path.join(self.base_dir, name)
      f = open(path, 'w')
      self.handles[name] = f
    return self.handles[name]

  def run(self):
    while True:
      item = self.in_channel.get()
      if item is SENTINEL:
        log('Disk thread stopped')
        break

      name, value = item
      f = self._file(name)

      if value is SENTINEL:
        f.close()
      else:
        f.write(value)


class QueueMultiplexer(object):
  """You can redirect multiple streams to this multiplexer with an ID.
  
  Then it will put them onto a queue/channel.
  """

  def __init__(self, out_channel):
    self.out_channel = out_channel 

  def redirect(self, stream, filename):
    """
    This should only be called from the event loop.
    """
    def on_data(chunk):
      self.out_channel.put((filename, chunk))

    def on_close():
      self.out_channel.put((filename, SENTINEL))

    stream.on('data', on_data)
    stream.on('close', on_close)


class Multiplexer(base.EventEmitter):
  """
  Combines multiple chunk streams into a single one.
  """
  EMITS = [
    'chunk',  # A tnet-encoded chunk of (stream ID, payload)
    'close',  # Emitted after ALL component streams are closed
    ]

  def __init__(self, event_loop=None):
    base.EventEmitter.__init__(self, event_loop=event_loop)
    self.all_closed = False
    # stream ID -> whether it was closed
    self.closed = {}

  def on_chunk(self, substream, chunk):
    log('MUX chunk %s %s', substream, chunk)
    # Need 2 different chunk values, since it imitates the TnetValueStream
    self.emit('chunk', tnet.dumps(substream))
    # Doesn't need to be serialized
    self.emit('chunk', chunk)

  def on_close(self, substream):
    """
    TODO: Emit stream ID on channel ZERO.  0 is reserved for close events.
    """
    log('MUX close %s', substream)
    self.closed[substream] = True
    if all(self.closed.values()):
      self.emit('close')

  def redirect(self, source, substream):
    self.closed[substream] = False
    assert isinstance(substream, int), "Got non-int substream %r" % substream

    # TODO: Should this behave like pipe(), with backpressure?
    source.on('chunk', lambda value: self.on_chunk(substream, value))
    source.on('close', lambda: self.on_close(substream))


class Demultiplexer(base.EventEmitter):
  """
  """
  def __init__(self, source, event_loop=None):
    base.EventEmitter.__init__(self, event_loop=event_loop)
    #self.source = source
    self.dests = {}
    source.on('chunk', self.on_chunk)
    source.on('close', self.on_close)
    self.last_id = None
    self.closed = {}

  def on_chunk(self, chunk):
    # This should wrap a TnetValueStream
    if self.last_id is None:
      # expecting an ID
      self.last_id = tnet.loads(chunk)
    else:
      dest = self.dests[self.last_id]
      dest.emit('chunk', chunk)
      self.last_id = None

    # TODO: chunk for Id == 0, and CLOSE

  def on_close(self):
    """
    """
    # if all the streams weren't closed already, close them now.
    for substream, closed in self.closed.iteritems():
      if not closed:
        self.dests[substream].emit('close')

  def redirect(self, substream, dest):
    self.dests[substream] = dest


class TnetOutput(base.EventEmitter):

  EMITS = ['chunk', 'close']

  def __init__(self, source, event_loop=None):
    base.EventEmitter.__init__(self, event_loop=event_loop)
    source.on('value', lambda value: self.emit('chunk', tnet.dumps(value)))
    source.on('close', lambda: self.emit('close'))
